import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import seaborn as sns
from datetime import datetime, timedelta
from scipy import signal
import warnings

warnings.filterwarnings('ignore')

# 设置中文字体和样式
plt.rcParams['font.sans-serif'] = ['SimHei', 'DejaVu Sans', 'Microsoft YaHei']
plt.rcParams['axes.unicode_minus'] = False
sns.set_style("whitegrid")
plt.rcParams['figure.figsize'] = [12, 8]


def enhance_holiday_features(input_file, output_file):
    """
    对原始数据进行处理，突出节假日特征并降低波动性
    """

    print(" 开始读取原始数据...")
    # 读取数据
    df = pd.read_csv(input_file)
    df['ds'] = pd.to_datetime(df['ds'])

    print(f" 原始数据统计:")
    print(f"   - 数据点数: {len(df)}")
    print(f"   - 平均值: {df['y'].mean():.2f}")
    print(f"   - 标准差: {df['y'].std():.2f}")
    print(f"   - 变异系数: {(df['y'].std() / df['y'].mean()) * 100:.2f}%")

    # 保存原始数据
    original_data = df.copy()

    # 第一步：平滑整体数据（降低波动性）
    print(" 进行数据平滑处理...")

    # 使用Savitzky-Golay滤波器进行平滑
    window_length = min(15, len(df) // 10 * 2 + 1)
    if window_length % 2 == 0:
        window_length += 1

    df['y_smoothed'] = signal.savgol_filter(df['y'], window_length, 2)

    # 第二步：识别并增强节假日特征
    print(" 增强节假日特征...")

    # 定义重要节假日
    holidays = {
        'childrens_day_2024': pd.Timestamp('2024-06-01'),
        'childrens_day_2025': pd.Timestamp('2025-06-01'),
        'christmas_2024': pd.Timestamp('2024-12-25'),
        'newyear_2025': pd.Timestamp('2025-01-01'),
        'thanksgiving_2024': pd.Timestamp('2024-11-28')
    }

    # 创建节假日增强因子
    df['holiday_boost'] = 1.0

    for holiday_name, holiday_date in holidays.items():
        # 设置节假日影响窗口
        if 'childrens_day' in holiday_name:
            # 六一儿童节：提前10天开始，持续到节后5天
            window_start = holiday_date - timedelta(days=10)
            window_end = holiday_date + timedelta(days=5)
            peak_boost = 2.5  # 六一儿童节增强2.5倍
        elif 'christmas' in holiday_name:
            # 圣诞节：提前20天开始，持续到节后10天
            window_start = holiday_date - timedelta(days=20)
            window_end = holiday_date + timedelta(days=10)
            peak_boost = 3.0  # 圣诞节增强3倍
        else:
            # 其他节日：提前7天开始，持续到节后3天
            window_start = holiday_date - timedelta(days=7)
            window_end = holiday_date + timedelta(days=3)
            peak_boost = 2.0

        # 应用节假日增强
        mask = (df['ds'] >= window_start) & (df['ds'] <= window_end)
        days_from_holiday = (df['ds'] - holiday_date).dt.days.abs()

        # 使用高斯函数创建平滑的增强曲线
        boost_values = peak_boost * np.exp(
            -(days_from_holiday ** 2) / (2 * (5 if 'childrens_day' in holiday_name else 8) ** 2))
        df.loc[mask, 'holiday_boost'] = np.maximum(df.loc[mask, 'holiday_boost'], boost_values[mask])

    # 第三步：应用节假日增强
    df['y_enhanced'] = df['y_smoothed'] * df['holiday_boost']

    # 第四步：后处理（确保数据合理性）
    df['y_enhanced'] = df['y_enhanced'].clip(lower=0)  # 确保非负
    df['y_enhanced'] = df['y_enhanced'].round().astype(int)  # 取整

    # 保持总体数据量基本不变
    total_original = df['y'].sum()
    total_enhanced = df['y_enhanced'].sum()
    if total_enhanced > 0:
        adjustment_factor = total_original / total_enhanced
        df['y_enhanced'] = (df['y_enhanced'] * adjustment_factor).round().astype(int)

    # 第五步：创建节假日标记特征（用于后续训练）
    df['is_childrens_day_period'] = 0
    df['is_christmas_period'] = 0

    for holiday_name, holiday_date in holidays.items():
        if 'childrens_day' in holiday_name:
            window_start = holiday_date - timedelta(days=10)
            window_end = holiday_date + timedelta(days=5)
            mask = (df['ds'] >= window_start) & (df['ds'] <= window_end)
            df.loc[mask, 'is_childrens_day_period'] = 1

        elif 'christmas' in holiday_name:
            window_start = holiday_date - timedelta(days=20)
            window_end = holiday_date + timedelta(days=10)
            mask = (df['ds'] >= window_start) & (df['ds'] <= window_end)
            df.loc[mask, 'is_christmas_period'] = 1

    print(" 数据处理完成!")
    print(f" 增强后数据统计:")
    print(f"   - 平均值: {df['y_enhanced'].mean():.2f}")
    print(f"   - 标准差: {df['y_enhanced'].std():.2f}")
    print(f"   - 变异系数: {(df['y_enhanced'].std() / df['y_enhanced'].mean()) * 100:.2f}%")

    # 保存结果
    output_df = df[['ds', 'y_enhanced', 'is_childrens_day_period', 'is_christmas_period']]
    output_df.columns = ['ds', 'y', 'is_childrens_day_period', 'is_christmas_period']
    output_df.to_csv(output_file, index=False)

    print(f" 增强后的数据已保存到: {output_file}")

    # 生成分析图表
    generate_analysis_charts(original_data, df, holidays)

    return df


def generate_analysis_charts(original_df, enhanced_df, holidays):
    """生成数据分析图表"""

    print(" 生成分析图表...")

    # 创建综合对比图
    fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, figsize=(16, 12))
    fig.suptitle('数据增强效果分析 - 节假日特征突出', fontsize=16, fontweight='bold')

    # 1. 原始数据 vs 增强数据
    ax1.plot(original_df['ds'], original_df['y'], 'b-', alpha=0.6, label='原始数据', linewidth=1)
    ax1.plot(enhanced_df['ds'], enhanced_df['y_enhanced'], 'r-', label='增强数据', linewidth=1.5)

    # 标记节假日
    for holiday_name, holiday_date in holidays.items():
        if 'childrens_day' in holiday_name:
            color = 'orange'
            label = '六一儿童节'
        elif 'christmas' in holiday_name:
            color = 'green'
            label = '圣诞节'
        else:
            color = 'purple'
            label = '其他节日'

        ax1.axvline(x=holiday_date, color=color, linestyle='--', alpha=0.7, label=label)

    ax1.set_title('数据对比: 原始 vs 增强', fontweight='bold')
    ax1.set_xlabel('日期')
    ax1.set_ylabel('出口数量')
    ax1.legend()
    ax1.grid(True, alpha=0.3)

    # 2. 节假日增强因子
    ax2.plot(enhanced_df['ds'], enhanced_df['holiday_boost'], 'g-', linewidth=2)
    ax2.axhline(y=1.0, color='gray', linestyle='--', alpha=0.5, label='基准线')

    for holiday_name, holiday_date in holidays.items():
        if 'childrens_day' in holiday_name:
            color = 'orange'
        elif 'christmas' in holiday_name:
            color = 'green'
        else:
            color = 'purple'

        ax2.axvline(x=holiday_date, color=color, linestyle='--', alpha=0.7)

    ax2.set_title('节假日增强因子时序图', fontweight='bold')
    ax2.set_xlabel('日期')
    ax2.set_ylabel('增强因子')
    ax2.legend()
    ax2.grid(True, alpha=0.3)
    ax2.set_ylim(0.5, 3.5)

    # 3. 六一儿童节特写
    childrens_day = pd.Timestamp('2024-06-01')
    window_start = childrens_day - timedelta(days=15)
    window_end = childrens_day + timedelta(days=10)

    mask = (enhanced_df['ds'] >= window_start) & (enhanced_df['ds'] <= window_end)
    zoom_data = enhanced_df[mask].copy()

    ax3.plot(zoom_data['ds'], original_df.loc[mask, 'y'], 'bo-', label='原始数据', markersize=4)
    ax3.plot(zoom_data['ds'], zoom_data['y_enhanced'], 'ro-', label='增强数据', markersize=4)
    ax3.axvline(x=childrens_day, color='orange', linestyle='--', linewidth=2, label='六一儿童节')
    ax3.set_title('六一儿童节特征增强特写', fontweight='bold')
    ax3.set_xlabel('日期')
    ax3.set_ylabel('出口数量')
    ax3.legend()
    ax3.grid(True, alpha=0.3)
    ax3.tick_params(axis='x', rotation=45)

    # 4. 圣诞节特写
    christmas = pd.Timestamp('2024-12-25')
    window_start = christmas - timedelta(days=25)
    window_end = christmas + timedelta(days=15)

    mask = (enhanced_df['ds'] >= window_start) & (enhanced_df['ds'] <= window_end)
    zoom_data = enhanced_df[mask].copy()

    ax4.plot(zoom_data['ds'], original_df.loc[mask, 'y'], 'bo-', label='原始数据', markersize=4)
    ax4.plot(zoom_data['ds'], zoom_data['y_enhanced'], 'ro-', label='增强数据', markersize=4)
    ax4.axvline(x=christmas, color='green', linestyle='--', linewidth=2, label='圣诞节')
    ax4.set_title('圣诞节特征增强特写', fontweight='bold')
    ax4.set_xlabel('日期')
    ax4.set_ylabel('出口数量')
    ax4.legend()
    ax4.grid(True, alpha=0.3)
    ax4.tick_params(axis='x', rotation=45)

    plt.tight_layout()
    plt.savefig('holiday_feature_enhancement_analysis.png', dpi=300, bbox_inches='tight')
    plt.show()

    # # 生成统计对比图
    # fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(14, 6))
    #
    # # 波动性对比
    # window = 7
    # original_volatility = original_df['y'].rolling(window=window).std()
    # enhanced_volatility = enhanced_df['y_enhanced'].rolling(window=window).std()
    #
    # ax1.plot(original_df['ds'], original_volatility, 'b-', label='原始数据波动性', alpha=0.7)
    # ax1.plot(enhanced_df['ds'], enhanced_volatility, 'r-', label='增强数据波动性', linewidth=1.5)
    # ax1.set_title('数据波动性对比 (7天移动标准差)', fontweight='bold')
    # ax1.set_xlabel('日期')
    # ax1.set_ylabel('标准差')
    # ax1.legend()
    # ax1.grid(True, alpha=0.3)
    #
    # # 分布对比
    # ax2.hist(original_df['y'], bins=30, alpha=0.6, label='原始数据', color='blue')
    # ax2.hist(enhanced_df['y_enhanced'], bins=30, alpha=0.6, label='增强数据', color='red')
    # ax2.set_title('数据分布对比', fontweight='bold')
    # ax2.set_xlabel('出口数量')
    # ax2.set_ylabel('频次')
    # ax2.legend()
    # ax2.grid(True, alpha=0.3)
    #
    # plt.tight_layout()
    # plt.savefig('data_statistics_comparison.png', dpi=300, bbox_inches='tight')
    # plt.show()

    # 生成处理报告
    generate_processing_report(original_df, enhanced_df)


def generate_processing_report(original_df, enhanced_df):
    """生成数据处理报告"""

    report = []
    report.append("=" * 60)
    report.append(" 数据增强处理报告")
    report.append("=" * 60)
    report.append(f"处理时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    report.append("")

    report.append(" 处理效果统计:")
    report.append(f"   - 原始数据平均值: {original_df['y'].mean():.2f}")
    report.append(f"   - 增强数据平均值: {enhanced_df['y_enhanced'].mean():.2f}")
    report.append(f"   - 原始数据标准差: {original_df['y'].std():.2f}")
    report.append(f"   - 增强数据标准差: {enhanced_df['y_enhanced'].std():.2f}")
    report.append(f"   - 原始变异系数: {(original_df['y'].std() / original_df['y'].mean()) * 100:.2f}%")
    report.append(
        f"   - 增强变异系数: {(enhanced_df['y_enhanced'].std() / enhanced_df['y_enhanced'].mean()) * 100:.2f}%")
    report.append("")

    report.append(" 节假日特征增强:")
    report.append("   - 六一儿童节: 提前10天开始，增强2.5倍")
    report.append("   - 圣诞节: 提前20天开始，增强3.0倍")
    report.append("   - 其他节日: 提前7天开始，增强2.0倍")
    report.append("")

    report.append(" 生成文件:")
    report.append("   - holiday_enhanced_data.csv (增强后的数据)")
    report.append("   - holiday_feature_enhancement_analysis.png (增强效果分析图)")
    # report.append("   - data_statistics_comparison.png (统计对比图)")
    report.append("")
    report.append("处理完成！增强后的数据更适合节假日特征学习")

    with open('data_enhancement_report.txt', 'w', encoding='utf-8') as f:
        f.write('\n'.join(report))

    print('\n'.join(report))
    print(" 处理报告已保存到: data_enhancement_report.txt")


# 主程序
if __name__ == "__main__":
    input_file = "prophet-data-simulated-2025-09-05-factor.csv"

    output_file = "holiday_enhanced_data_not_smoothed.csv"

    print(" 开始数据增强处理...")
    print("=" * 60)

    enhanced_data = enhance_holiday_features(input_file, output_file)

    print("=" * 60)
    print(" 数据增强处理完成！")
    print("=" * 60)
    print("下一步建议:")
    print("使用增强后的数据进行Prophet模型训练，节假日特征会更加明显")
    print("在模型中可以添加 is_childrens_day_period 和 is_christmas_period 作为回归器")